"""
决策树分类：基于信息增益或基尼指数
核心思想：选择最佳特征进行划分，使得划分后的数据尽可能地纯净
信息增益：父熵 - 子熵  -> 越大越好 ， 侧重方向：信息量
基尼指数：1 - 类别概率平方和  -> 越小越好 ， 侧重方向：类别分布
"""
import pandas as pd    # titanic数据集是csv文件，所以用pandas读取
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.feature_extraction import DictVectorizer
from sklearn.tree import DecisionTreeClassifier, export_graphviz



# 读取数据
data = pd.read_csv("../src/titanic/titanic.csv")

# 筛选目标值和特征值
x = data[["pclass", "age", "sex"]]
y = data["survived"]

# 缺失值处理
x["age"] = x["age"].fillna(x["age"].mean())

# 转换为字典
x = x.to_dict(orient="records")
"""
用pandas读取csv文件，转换成DataFrame格式
DataFrame格式的数据，不能直接进行训练，需要转换成字典格式，进行字典特征提取
"""

# 数据集划分
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=22)
"""
使用网格搜索，也不能因为里面带有K折交叉验证，误认为会自动划分数据集
实质：train_test_split划分出训练集和测试集
     网格搜索会在训练集上进行交叉验证，得到最优参数
     然后再用最优参数在测试集上进行测试
"""

# 字典特征提取
transfer = DictVectorizer()
x_train = transfer.fit_transform(x_train)
x_test = transfer.transform(x_test)

# 创建决策树分类器
estimator = DecisionTreeClassifier(random_state=22)

# 创建网格搜索对象
param_grid = {"criterion": ["gini", "entropy"],
              "max_depth": [3, 5, 7, 9, 11]}
model = GridSearchCV(estimator, param_grid=param_grid, cv=5)

# 模型训练
model.fit(x_train, y_train)

# 查看结果
print("最佳参数：", model.best_params_)
print("最佳分数：", model.best_score_)

# 模型评估 -> 上面都是训练集上的结果，模型评估是在测试集上的结果
print("测试集准确率：", model.score(x_test, y_test))

# 可视化决策树
export_graphviz(model.best_estimator_,
                out_file="titanic_tree.dot",
                feature_names=transfer.get_feature_names_out())
